package EP::Visualization::Collapser;
use EP::Visualization::Hooks;
use EP::Visualization::HeatmapGraphic;
use EP::Visualization::TreeGraphic;
use strict;


sub new {
  my $proto = shift;
  my $class = ref($proto) || $proto;
  my $self = {};
  bless $self, $class;
  return $self;
}

#############################################################
#Determines graphics objects present and collapses these    #
#graphic objects by calling their specific collapse methods.#
#############################################################

sub collapse {
  my $self = shift;
  my $size = shift;
  my $type = shift;              #type determines type of collapsing, either collapsing rows or columns
  my @objects  = @_;
  my $treeIndex = -1;            #will contain index of tree in the objects array
  my $heatmapIndex = -1;         #will contain index of heatmap in the objects array
  my $annotIndex = -1;           #will contain index of annotation in the objects array
  my %revIndex;

  for my $i ( 0 .. $#objects ) {
    $treeIndex    = $i if $objects[$i]->getType() eq "tree";
    $heatmapIndex = $i if $objects[$i]->getType() eq "heatmap";
    $annotIndex   = $i if $objects[$i]->getType() eq "annotation";
  }

  die "No tree to collapse to" if ($treeIndex == -1);
  die "Nothing to collapse" if ($heatmapIndex == -1 and $annotIndex == -1 );
  %revIndex   =  collapseTree ( $objects[$treeIndex], $size );

  if ( $heatmapIndex != -1 ) {
#    collapseHeatmap ( $objects[$heatmapIndex], $objects[$treeIndex], \%revIndex, $type );
  }

  if ( $annotIndex != -1 ) {
    collapseAnnotation ( $objects[$annotIndex], $objects[$treeIndex], \%revIndex );
  }
}

#########################################################
#TreeGraphic specific collapse method, reduces tree to  #
#size (percentage of original) that should remain       #
#########################################################

sub collapseTree {
  my $tree = shift;
  my $size = shift;
  my @treeData = @{ $tree -> getData()};
  my @newTreeData;
  my $threshold = int((( $#treeData + 1 )/100) * $size);  #determines threshold index of data
  my $index = 0;
  my $newIndex = $#treeData + 10;  #the new index may not exist yet in the old hookset so we add one to the highest original hook
  my %revIndex;                 #revIndex maps a new index to an array which contains the old hooks that are mapped to this new index
  my %newIndices;                #hash that goes from old_index -> new_index where the old_index is the original hookset and the new_index the collapsed hookset

  for my $y ( 0..$#treeData - $threshold ){    #next we go from the buttom of the treedata up to the threshold index any nodes below the index will be replaced by new indices#
    my $i = $#treeData - $y;
    if (!exists $newIndices{ $treeData[ $i ][ 0 ] } && !exists $newIndices{ $treeData[ $i ][ 1 ] }){  #both indices haven't got a new index yet so they both get the same new index
      $newIndices{ $treeData[ $i ][ 0 ] } = $newIndex;
      $newIndices{ $treeData[ $i ][ 1 ] } = $newIndex;
      $revIndex{ $newIndex } = ([ $treeData[ $i ][ 0 ], $treeData[ $i ][ 1 ] ]);
      $newIndex++;
    }
    
    elsif ( !exists $newIndices{ $treeData[ $i ][ 0 ] }){ #means one of the two indices was already mapped to a new index so the one that was't should get this new index too
      my $newInd = $newIndices{ $treeData[ $i ][ 1 ] };
      $newIndices{ $treeData[ $i ][ 0 ] } = $newInd;
      my @prev = @{ $revIndex{ $newInd } };
      $prev[ $#prev+1 ]= $treeData[ $i ][ 0 ];
      $revIndex{ $newInd } = \@prev;
    }
    
    elsif(!exists $newIndices{ $treeData[ $i ][ 1 ] }){  #means one of the two indices was already mapped to a new index so the one that was't should get this new index too
      my $newInd =$newIndices{ $treeData[ $i ][ 0 ] };
      $newIndices{ $treeData[ $i ][ 1 ] } = $newInd;
      my @prev = @{ $revIndex{ $newInd } };
      $prev[ $#prev+1 ] = $treeData[ $i ][ 1 ];
      $revIndex{ $newInd } = \@prev;
    }
    
    elsif ( $newIndices{ $treeData[ $i ][ 0 ] } != $newIndices{ $treeData[ $i ][ 1 ] }){ #Both indices have already a new index but a different one, so we should eliminate one of     
      my $Ind0 = $newIndices{ $treeData[ $i ][ 0 ] };                                #the new indices and map to the same new index
      my $Ind1 = $newIndices{ $treeData[ $i ][ 1 ] };
      $newIndices{ $treeData[ $i ][ 1 ]} = $Ind0;
      my @prev1 = @{ $revIndex{ $Ind0 } };
      my @prev2 = @{ $revIndex{ $Ind1 } };
      for my $r ( 0..$#prev2 ){
        $prev1[ $#prev1+1 ] = $prev2[$r];
      }
      $revIndex{ $Ind0 } = \@prev1 ;
    }
  }
  for ( my $i=0 ; $i < $threshold ; $i++ ){                         #next old indices that are lower than the threshold are replaced by their new indices if those exist   
    if ( exists $newIndices{ $treeData[ $i ][ 0 ] } && exists $newIndices{ $treeData[ $i ][ 1 ] }){
      $newTreeData[ $i ] = ([ $newIndices{ $treeData[ $i ][ 0 ] }, $newIndices{ $treeData[ $i ][ 1 ] }, $treeData[ $i ][ 2 ], $treeData[ $i ][ 3 ], $treeData[ $i ][ 4 ] ]);
    }
    elsif ( exists $newIndices{ $treeData[ $i ][ 1 ] }){
      $newTreeData[ $i ] = ([ $treeData[ $i ][ 0 ],$newIndices{ $treeData[ $i ][ 1 ]}, $treeData[ $i ][ 2 ], $treeData[ $i ][ 3 ],$treeData[ $i ][ 4 ]]);
    }
    elsif ( exists $newIndices{ $treeData[ $i ][ 0 ] }){
      $newTreeData[ $i ] = ([ $newIndices{ $treeData[ $i ][ 0 ] }, $treeData[ $i ][ 1 ], $treeData[ $i ][ 2 ], $treeData[ $i ][ 3 ], $treeData[ $i ][ 4 ]]);
    }
    else{
      $newTreeData[ $i ] = $treeData[ $i ];
    }
  }
  
  $tree -> update( \@newTreeData );  
  return %revIndex;
}


 
 ########################################################################################
 #next the heatmap data are collapsed according to the collapsed tree                   # 
 #                                                                                      #
 #starting from the new treehooks we fill in the heatmap datamatrix                     # 
 #if we encounter a new hook index then we take all the values that were present in the #
 #old datamatrix corresponding to those hooks and average them                          #
 ######################################################################################## 

sub collapseHeatmap{
  
  my $heatmap = shift;
  my $tree = shift;
  my $revInd = shift;
  my $type = shift;
  my %revIndex = %{ $revInd };

  if ($type eq "row"){
    my @newHeatmapData;
    my @heatmapData = @{ $heatmap -> getData()};
    my @heatmapHook = @{ $heatmap -> getVerticalHookSet -> getHookSet()};
    my @treeHook = @{ $tree -> getHook -> getHookSet()};
    my %hook;
    
    for my $j ( 0..$#heatmapHook ){
      $hook{ $heatmapHook[ $j ] } = $j;      #is used as a fast way to go from hook to position in heatmap array
    }
    
    for my $i ( 0..$#treeHook ){
      if( exists $revIndex{ $treeHook[ $i ] } ){
	my @hooks = @{ $revIndex{ $treeHook[ $i ] } };
	my @newRow = @{ $heatmapData[ $hook{ $hooks[ 0 ] - 1 } ]};
	for my $d ( 1..$#hooks ){
	  my @old = @{ $heatmapData[ $hook{ $hooks[ $d ] - 1 } ] };
	  for my $u ( 0..$#old ){
	    $newRow[ $u ] = ($newRow[ $u ]+ $old[ $u ])/2 ;
	  }
	}
	push @newHeatmapData, \@newRow;
      }
      else{
	push @newHeatmapData, $heatmapData[ $hook{ $treeHook[ $i ] } ]; 
      }
    }
    
    $heatmap -> setData( data => \@newHeatmapData );
    $heatmap -> setVerticalHookSet( $tree -> getHook());
  }
  
  elsif ($type eq "column"){
    my @newHeatmapData = ();
    my @heatmapData = @{ $heatmap  -> getData()};
    my @heatmapHook = @{ $heatmap  -> getHorizontalHookSet -> getHookSet()};
    my @treeHook = @{ $tree  -> getHook -> getHookSet()};
    my %hook;
    my $rows = $#heatmapData;
    
    for my $j ( 0..$#heatmapHook ){
      $hook{ $heatmapHook[ $j ] } = $j;                      #is used as a fast way to go from hook to position in heatmap array
    }
    
    for my $y ( 0..$#heatmapData ){
      my @newRow = ();
      my @oldRow = @{ $heatmapData[ $y ] };
      
      for my $i ( 0..$#treeHook ){
	if( exists $revIndex{ $treeHook[ $i ] + 1 }){
	  my @hooks = @{ $revIndex{ $treeHook[ $i ] + 1 } };
	  my $value = $oldRow[ $hook{ $hooks[ 0 ] - 1 } ];
	  
	  for my $h ( 1..$#hooks ){
	    $value = ($value + $oldRow[ $hook{ $hooks[ $h ] - 1 } ] )/2;  
	  }	 
	  push @newRow, $value;
	} 
	else{
	  push @newRow, $oldRow[ $hook{ $treeHook[ $i ] } ];
	}
      }
      push @newHeatmapData, \@newRow;
    }
    $heatmap -> setData( \@newHeatmapData );
    $heatmap  -> setHorizontalHookSet( $tree -> getHook());
    
  }
  else{ die "Non valid collapsing type"; }
}
########################################################################################
#Next we do the annotation collapsing if there is any annotation graphicsobject present#
########################################################################################


sub collapseAnnotation{
  
  my $annotation = shift;
  my $tree = shift;
  my $revInd = shift;
  my %revIndex = %{ $revInd };
  my @treeHook = @{ $tree -> getHook -> getHookSet()};
  my @annotHook = @{ $annotation -> getHook -> getHookSet()};
  my @annotData = @{ $annotation  -> getData()};
  my @newAnnotData = ();
  my %hook;
  
  for my $j ( 0..$#annotHook ){
    $hook{ $annotHook[ $j ] } = $j;                      #is used as a fast way to go from hook to position in heatmap array
  }
  
  for my $i ( 0..$#treeHook ){
    if( exists $revIndex{ $treeHook[ $i ] + 1 } ){           #if there is a collapsed branch we'll look up how many leaves are beneath it and use this as new annotation
      my $numOfLeaves = scalar( @{ $revIndex{ $treeHook[ $i ] + 1 } });
      push @newAnnotData, "[$numOfLeaves]"; 
    }
    else{
      push @newAnnotData, $annotData[ $hook{ $treeHook[ $i ] } ];  #here we use the original annotation as this is an uncollapsed leaf
    }
  }
  $annotation -> setData( data => \@newAnnotData );
  $annotation -> setHook( $tree  -> getHook());
}  
1;
